# -*- coding: utf-8 -*-
r"""
NCL
################################################
Reference:
    Zihan Lin*, Changxin Tian*, Yupeng Hou*, Wayne Xin Zhao. "Improving Graph Collaborative Filtering with Neighborhood-enriched Contrastive Learning." in WWW 2022.
"""

import faiss
import numpy as np
import scipy.sparse as sp
import torch
import torch.nn.functional as F
from recbole.model.abstract_recommender import GeneralRecommender
from recbole.model.init import xavier_uniform_initialization
from recbole.model.loss import BPRLoss, EmbLoss
from recbole.utils import InputType
from reckit import randint_choice
from graph_Augmentation.aug_main import inter_random_augment


# 模型在训练时靠self.calculate_loss()串起流程
class NCL(GeneralRecommender):
    input_type = InputType.PAIRWISE

    def __init__(self, config, dataset):
        super(NCL, self).__init__(config, dataset)

        # 拿到交互矩阵
        # 这里的交互矩阵是scipy.sparse._coo创建的矩阵
        # 可以通过.tocsr()转换成坐标+值的稀疏矩阵形式
        self.interaction_matrix = dataset.inter_matrix(form='coo').astype(np.float32)

        # 加载参数
        self.latent_dim = config['embedding_size']  # int type: the embedding size of the base model
        self.n_layers = config['n_layers']  # int type: the layer num of the base model
        self.reg_weight = config['reg_weight']  # float32 type: the weight decay for l2 normalization

        self.ssl_temp = config['ssl_temp']
        self.ssl_reg = config['ssl_reg']
        self.hyper_layers = config['hyper_layers']

        self.alpha = config['alpha']

        self.proto_reg = config['proto_reg']
        self.k = config['num_clusters']

        # 模型初始化
        self.user_embedding = torch.nn.Embedding(num_embeddings=self.n_users, embedding_dim=self.latent_dim)
        self.item_embedding = torch.nn.Embedding(num_embeddings=self.n_items, embedding_dim=self.latent_dim)

        self.mf_loss = BPRLoss()
        self.reg_loss = EmbLoss()

        # storage variables for full sort evaluation acceleration
        # 为了全排序评估加速设置的变量
        self.restore_user_e = None
        self.restore_item_e = None

        # 基于随机边扰动数据增强产生增强交互图的邻接矩阵
        # 这个是扰动去除边的比例
        self.leave_ratio = 0.2
        # 指定随机增强加边的参数比例
        self.density_sample_ratio = config['density_sample_ratio']
        # print("density_sample_ratio", self.density_sample_ratio)
        # self.aug_inter_matrix1 = self.inter_matrix_augment()
        # self.aug_inter_matrix2 = self.inter_matrix_augment()
        self.aug_inter_matrix1 = self.inter_matrix_add_augment()
        self.aug_inter_matrix2 = self.inter_matrix_add_augment()
        print("aug_matrix1", self.aug_inter_matrix1.shape)
        print("aug_matrix2", self.aug_inter_matrix2.shape)

        # 进行交互矩阵的规范化
        self.norm_adj_mat = self.get_norm_adj_mat(self.interaction_matrix).to(self.device)
        # 进行两个增强交互矩阵的规范化
        self.aug_norm_mat1 = self.get_norm_adj_mat(self.aug_inter_matrix1).to(self.device)
        self.aug_norm_mat2 = self.get_norm_adj_mat(self.aug_inter_matrix2).to(self.device)


        # print("hahah:", self.aug_inter_matrix1.shape)
        # print("hahah:", self.interaction_matrix.shape)

        # 实现模型参数随机初始化
        self.apply(xavier_uniform_initialization)
        self.other_parameter_name = ['restore_user_e', 'restore_item_e']

        self.user_centroids = None
        self.user_2cluster = None
        self.item_centroids = None
        self.item_2cluster = None

        # 原图卷积聚类映射向量
        self.user_conv_centroids = None
        self.user_conv_2cluster = None
        self.item_conv_centroids = None
        self.item_conv_2cluster = None

        # 增强图1卷积聚类映射向量
        self.aug_user_conv_centroids1 = None
        self.aug_user_conv_2cluster1 = None
        self.aug_item_conv_centroids1 = None
        self.aug_item_conv_2cluster1 = None

        # 增强图2卷积聚类映射向量
        self.aug_user_conv_centroids2 = None
        self.aug_user_conv_2cluster2 = None
        self.aug_item_conv_centroids2 = None
        self.aug_item_conv_2cluster2 = None

    def update_aug_conv_kmeans(self):
        # 对增强图进行的卷积聚类
        aug_user_all_emb1, aug_item_all_emb1, _ = self.forward(self.aug_norm_mat1)
        aug_user_conv_embeddings1 = aug_user_all_emb1.detach().cpu().numpy()
        aug_item_conv_embeddings1 = aug_item_all_emb1.detach().cpu().numpy()
        self.aug_user_conv_centroids1, self.aug_user_conv_2cluster1 = self.run_kmeans(aug_user_conv_embeddings1)
        self.aug_item_conv_centroids1, self.aug_item_conv_2cluster1 = self.run_kmeans(aug_item_conv_embeddings1)

        aug_user_all_emb2, aug_item_all_emb2, _ = self.forward(self.aug_norm_mat2)
        aug_user_conv_embeddings2 = aug_user_all_emb2.detach().cpu().numpy()
        aug_item_conv_embeddings2 = aug_item_all_emb2.detach().cpu().numpy()
        self.aug_user_conv_centroids2, self.aug_user_conv_2cluster2 = self.run_kmeans(aug_user_conv_embeddings2)
        self.aug_item_conv_centroids2, self.aug_item_conv_2cluster2 = self.run_kmeans(aug_item_conv_embeddings2)

    def update_conv_kmeans(self):
        # 获取当前用户项目表征向量的卷积
        user_conv_embeddings, item_conv_embeddings, _ = self.forward(self.norm_adj_mat)
        # 将向量由tensor转为numpy数据格式
        user_conv_embeddings = user_conv_embeddings.detach().cpu().numpy()
        item_conv_embeddings = item_conv_embeddings.detach().cpu().numpy()
        # 对卷积后的用户和项目分别进行K聚类
        self.user_conv_centroids, self.user_conv_2cluster = self.run_kmeans(user_conv_embeddings)
        self.item_conv_centroids, self.item_conv_2cluster = self.run_kmeans(item_conv_embeddings)

    def update_kmeans(self):
        # 获取当前用户项目表征向量
        user_embeddings = self.user_embedding.weight.detach().cpu().numpy()
        item_embeddings = self.item_embedding.weight.detach().cpu().numpy()
        # 对用户和项目分别进行K聚类
        self.user_centroids, self.user_2cluster = self.run_kmeans(user_embeddings)
        self.item_centroids, self.item_2cluster = self.run_kmeans(item_embeddings)

    def run_kmeans(self, x):
        # print("original x:", x.shape)
        """Run K-means algorithm to get k clusters of the input tensor x
        返回结果有二：
        centroids：对x进行k聚类得到的中心点坐标数组
        node2cluster：x中每个向量进行近邻映射后，对应的聚类中心索引，即centroids[node2cluster]
        """
        # 电脑没有GPU，就指定参数为false
        # 默认是1000中心的聚类，空间为64维
        kmeans = faiss.Kmeans(d=self.latent_dim, k=self.k, gpu=False)
        # 执行聚类
        kmeans.train(x)
        # 获得聚类的中心点坐标
        cluster_cents = kmeans.centroids
        # print("cluster_cents:", cluster_cents.shape)
        # 把待聚类的向量映射至最近的中心点,I保存了每个向量对应类中心的索引
        # 第一个参数是映射后的L2距离，用不上
        _, I = kmeans.index.search(x, 1)
        # print("mapped I:", I.shape)

        # convert to cuda Tensors for broadcast
        # 转换为tensor并规范化，方便广播
        centroids = torch.Tensor(cluster_cents).to(self.device)
        centroids = F.normalize(centroids, p=2, dim=1)

        # 投影索引转换成tensor并返回
        node2cluster = torch.LongTensor(I).squeeze().to(self.device)
        # print("node2cluster", node2cluster.shape)
        return centroids, node2cluster

    def get_norm_adj_mat(self, inter_matrix):
        r"""Get the normalized interaction matrix of users and items.

        Construct the square matrix from the training data and normalize it
        using the laplace matrix.
        基于拉普拉斯变换实现训练数据的规范化

        .. math::
            A_{hat} = D^{-0.5} \times A \times D^{-0.5}

        Returns:
            Sparse tensor of the normalized interaction matrix.
        """
        # 图卷积拉氏变换的专用大矩阵
        # 这是个字典，也就是坐标加值
        A = sp.dok_matrix((self.n_users + self.n_items, self.n_users + self.n_items), dtype=np.float32)
        # print("A", self.n_users * self.n_items)

        inter_M = inter_matrix
        # print("inter_m:\n", inter_M.shape, type(inter_M))
        inter_M_t = inter_M.transpose()
        # nnz，矩阵中的非零元素向量
        # 这里更新了左下和右上两个子矩阵
        # 第一个参数是坐标
        # 第二个参数是对应的值，需要可迭代的格式，所以乘以[1]成为list
        data_dict = dict(zip(zip(inter_M.row, inter_M.col + self.n_users), [1] * inter_M.nnz))
        # print("inter_M.row:", type(inter_M.row))
        # print("inter_M.col:", type(inter_M.col))
        # print("inter.nnz:", len(([1] * inter_M.nnz)))
        # print("data_dict", len(data_dict))
        # N
        data_dict.update(dict(zip(zip(inter_M_t.row + self.n_users, inter_M_t.col), [1] * inter_M_t.nnz)))
        # print("data_dict", len(data_dict))
        # 2N
        A._update(data_dict)
        # 给用户和项目对应的交互数求和
        sumArr = (A > 0).sum(axis=1)
        # print("sumArr", sumArr.shape)
        # (n_u+n_v, 1)
        # 后续要做除法，把零元素变成极小数
        # 向量拍扁然后对角化，得到D
        diag = np.array(sumArr.flatten())[0] + 1e-7
        diag = np.power(diag, -0.5)
        self.diag = torch.from_numpy(diag).to(self.device)
        # 从tensor构建一个sparseMatrix
        D = sp.diags(diag)
        # 经典的拉式分解变换
        L = D @ A @ D
        # 变成坐标模式，方便提取坐标和值，还原成tensor
        L = sp.coo_matrix(L)
        row = L.row
        col = L.col
        index = torch.LongTensor(np.array([row, col]))
        data = torch.FloatTensor(L.data)
        # 还原出拉氏变换后的矩阵
        SparseL = torch.sparse.FloatTensor(index, data, torch.Size(L.shape))
        return SparseL

    def inter_matrix_add_augment(self):
        """
        基于原有数据集的交互矩阵，进行增加类型的边扰动变换，产生新的增强图交互矩阵

        :return: aug_int_mat
        """
        aug_int_mat = inter_random_augment(self.interaction_matrix, self.density_sample_ratio)
        return aug_int_mat

    def inter_matrix_augment(self):
        """
        基于原有数据集的交互矩阵，进行边扰动变换，产生新的增强图交互矩阵

        :return: aug_int_mat
        """
        user_ids = self.interaction_matrix.row
        item_ids = self.interaction_matrix.col
        # 借用SGL的交互图随机增强函数，对坐标ID进行操作
        # 看起来是对user_ids进行操作，其实因为这里的用户和项目id都是未去重的，等于同时操作了user和item
        keep_idx = randint_choice(len(user_ids), size=int(len(user_ids) * (1 - self.leave_ratio)), replace=False)
        user_np = np.array(user_ids)[keep_idx]
        item_np = np.array(item_ids)[keep_idx]
        ratings = np.ones_like(user_np, dtype=np.float32)
        # 再生成稀疏图并返回
        aug_int_mat = sp.coo_matrix((ratings, (user_np, item_np)),
                                    shape=(self.n_users, self.n_items))
        # tmp_adj = sp.csr_matrix((ratings, (user_np, item_np + self.n_users)), shape=(n_nodes, n_nodes))

        return aug_int_mat

    def get_ego_embeddings(self):
        """Get the embedding of users and items and combine to an embedding matrix.
        拼接返回当前权重值
        Returns:
            Tensor of the embedding matrix. Shape of [n_items+n_users, embedding_dim]
        """
        user_embeddings = self.user_embedding.weight
        item_embeddings = self.item_embedding.weight
        ego_embeddings = torch.cat([user_embeddings, item_embeddings], dim=0)
        return ego_embeddings

    def forward(self, norm_adj_mat):
        all_embeddings = self.get_ego_embeddings()
        embeddings_list = [all_embeddings]
        # 根据设定的参数卷起来
        for layer_idx in range(max(self.n_layers, self.hyper_layers * 2)):
            all_embeddings = torch.sparse.mm(norm_adj_mat, all_embeddings)
            embeddings_list.append(all_embeddings)
        # 把卷积向量和初始向量放一起，求平均，得到用户和项目的更新向量
        lightgcn_all_embeddings = torch.stack(embeddings_list[:self.n_layers + 1], dim=1)
        lightgcn_all_embeddings = torch.mean(lightgcn_all_embeddings, dim=1)

        # 切分向量，连通全部卷积结果返回
        user_all_embeddings, item_all_embeddings = torch.split(lightgcn_all_embeddings, [self.n_users, self.n_items])
        return user_all_embeddings, item_all_embeddings, embeddings_list

    # 计算增强图卷积后向量基于语义聚类的对比损失
    def AugConvNCE_loss(self, aug_user_conv_embeddings_all, aug_item_conv_embeddings_all, user, item, aug_order):
        if aug_order == 1:
            aug_user_conv_2cluster = self.aug_user_conv_2cluster1
            aug_user_conv_centroids = self.aug_user_conv_centroids1
            aug_item_conv_2cluster = self.aug_item_conv_2cluster1
            aug_item_conv_centroids = self.aug_item_conv_centroids1
        # elif aug_order == 2:
        else:
            aug_user_conv_2cluster = self.aug_user_conv_2cluster2
            aug_user_conv_centroids = self.aug_user_conv_centroids2
            aug_item_conv_2cluster = self.aug_item_conv_2cluster2
            aug_item_conv_centroids = self.aug_item_conv_centroids2
        # assert aug_order in range(1, 2)
        # aug_user_conv_embeddings_all, aug_item_conv_embeddings_all = torch.split(node_embedding,
        #                                                                         [self.n_users, self.n_items])

        # 取出用户的向量并标准化
        aug_user_conv_embeddings = aug_user_conv_embeddings_all[user]  # [B, e]
        norm_aug_user_conv_embeddings = F.normalize(aug_user_conv_embeddings)
        # print("norm_aug_user_conv_embeddings", norm_aug_user_conv_embeddings.shape)
        # 获取用户向量对应的聚类中心索引
        user2cluster = aug_user_conv_2cluster[user]  # [B,]
        # 拿到向量映射的中心向量
        user2centroids = aug_user_conv_centroids[user2cluster]  # [B, e]
        # print("user2centroids", user2centroids.shape)
        # 以下开始计算对比损失的分子，正样本得分
        # 进行了逐元素相乘，以聚类中心向量对用户向量进行修正
        pos_score_user = torch.mul(norm_aug_user_conv_embeddings, user2centroids).sum(dim=1)
        # torch.exp(x)是e^x
        # ssl_temp是一个小数，疑似用于放大
        pos_score_user = torch.exp(pos_score_user / self.ssl_temp)
        # print("user_centroids", self.user_centroids.shape)
        # 计算对比损失的分母
        # 基于（total）全部中心向量得到一组用户向量评分
        ttl_score_user = torch.matmul(norm_aug_user_conv_embeddings, aug_user_conv_centroids.transpose(0, 1))
        ttl_score_user = torch.exp(ttl_score_user / self.ssl_temp).sum(dim=1)

        # 对比学习的总损失函数，即用户和自己映射的中心向量尽量接近，而和其他尽量离远些
        # ttl没有把用户映射的中心向量分出来，可能问题不大
        aug_conv_nce_loss_user = -torch.log(pos_score_user / ttl_score_user).sum()

        # 类似的项目向量对比学习损失函数计算
        aug_item_conv_embeddings = aug_item_conv_embeddings_all[item]
        norm_aug_conv_item_embeddings = F.normalize(aug_item_conv_embeddings)

        item2cluster = aug_item_conv_2cluster[item]  # [B, ]
        item2centroids = aug_item_conv_centroids[item2cluster]  # [B, e]
        pos_score_item = torch.mul(norm_aug_conv_item_embeddings, item2centroids).sum(dim=1)
        pos_score_item = torch.exp(pos_score_item / self.ssl_temp)
        ttl_score_item = torch.matmul(norm_aug_conv_item_embeddings, aug_item_conv_centroids.transpose(0, 1))
        ttl_score_item = torch.exp(ttl_score_item / self.ssl_temp).sum(dim=1)
        aug_conv_nce_loss_item = -torch.log(pos_score_item / ttl_score_item).sum()

        # 两个损失函数的加权和，权重是一个很小的数
        conv_nce_loss = self.proto_reg * (aug_conv_nce_loss_user + aug_conv_nce_loss_item)
        return conv_nce_loss

    # 计算原图卷积后向量基于语义聚类的对比损失
    def ConvNCE_loss(self, node_embedding, user, item):
        user_conv_embeddings_all, item_conv_embeddings_all = torch.split(node_embedding, [self.n_users, self.n_items])

        # 取出用户的向量并标准化
        user_conv_embeddings = user_conv_embeddings_all[user]  # [B, e]
        norm_user_conv_embeddings = F.normalize(user_conv_embeddings)
        # print("norm_user_conv_embeddings", norm_user_conv_embeddings.shape)
        # 获取用户向量对应的聚类中心索引
        user2cluster = self.user_conv_2cluster[user]  # [B,]
        # 拿到向量映射的中心向量
        user2centroids = self.user_conv_centroids[user2cluster]  # [B, e]
        # print("user2centroids", user2centroids.shape)
        # 以下开始计算对比损失的分子，正样本得分
        # 进行了逐元素相乘，以聚类中心向量对用户向量进行修正
        pos_score_user = torch.mul(norm_user_conv_embeddings, user2centroids).sum(dim=1)
        # torch.exp(x)是e^x
        # ssl_temp是一个小数，疑似用于放大
        pos_score_user = torch.exp(pos_score_user / self.ssl_temp)
        # print("user_centroids", self.user_centroids.shape)
        # 计算对比损失的分母
        # 基于（total）全部中心向量得到一组用户向量评分
        ttl_score_user = torch.matmul(norm_user_conv_embeddings, self.user_conv_centroids.transpose(0, 1))
        ttl_score_user = torch.exp(ttl_score_user / self.ssl_temp).sum(dim=1)

        # 对比学习的总损失函数，即用户和自己映射的中心向量尽量接近，而和其他尽量离远些
        # ttl没有把用户映射的中心向量分出来，可能问题不大
        conv_nce_loss_user = -torch.log(pos_score_user / ttl_score_user).sum()

        # 类似的项目向量对比学习损失函数计算
        item_conv_embeddings = item_conv_embeddings_all[item]
        norm_conv_item_embeddings = F.normalize(item_conv_embeddings)

        item2cluster = self.item_conv_2cluster[item]  # [B, ]
        item2centroids = self.item_conv_centroids[item2cluster]  # [B, e]
        pos_score_item = torch.mul(norm_conv_item_embeddings, item2centroids).sum(dim=1)
        pos_score_item = torch.exp(pos_score_item / self.ssl_temp)
        ttl_score_item = torch.matmul(norm_conv_item_embeddings, self.item_conv_centroids.transpose(0, 1))
        ttl_score_item = torch.exp(ttl_score_item / self.ssl_temp).sum(dim=1)
        conv_nce_loss_item = -torch.log(pos_score_item / ttl_score_item).sum()

        # 两个损失函数的加权和，权重是一个很小的数
        conv_nce_loss = self.proto_reg * (conv_nce_loss_user + conv_nce_loss_item)
        return conv_nce_loss

    # 计算卷积前向量基于语义聚类的对比损失
    def ProtoNCE_loss(self, node_embedding, user, item):
        user_embeddings_all, item_embeddings_all = torch.split(node_embedding, [self.n_users, self.n_items])
        # 取出用户的向量并标准化
        user_embeddings = user_embeddings_all[user]  # [B, e]
        norm_user_embeddings = F.normalize(user_embeddings)
        # print("norm_user_embeddings", norm_user_embeddings.shape)
        # 获取用户向量对应的聚类中心索引
        user2cluster = self.user_2cluster[user]  # [B,]
        # 拿到向量映射的中心向量
        user2centroids = self.user_centroids[user2cluster]  # [B, e]
        # print("user2centroids", user2centroids.shape)
        # 以下开始计算对比损失的分子，正样本得分
        # 进行了逐元素相乘，以聚类中心向量对用户向量进行修正
        pos_score_user = torch.mul(norm_user_embeddings, user2centroids).sum(dim=1)
        # torch.exp(x)是e^x
        # ssl_temp是一个小数，疑似用于放大
        pos_score_user = torch.exp(pos_score_user / self.ssl_temp)
        # print("user_centroids", self.user_centroids.shape)
        # 计算对比损失的分母
        # 基于（total）全部中心向量得到一组用户向量评分
        ttl_score_user = torch.matmul(norm_user_embeddings, self.user_centroids.transpose(0, 1))
        ttl_score_user = torch.exp(ttl_score_user / self.ssl_temp).sum(dim=1)

        # 对比学习的总损失函数，即用户和自己映射的中心向量尽量接近，而和其他尽量离远些
        # ttl没有把用户映射的中心向量分出来，可能问题不大
        proto_nce_loss_user = -torch.log(pos_score_user / ttl_score_user).sum()

        # 类似的项目向量对比学习损失函数计算
        item_embeddings = item_embeddings_all[item]
        norm_item_embeddings = F.normalize(item_embeddings)

        item2cluster = self.item_2cluster[item]  # [B, ]
        item2centroids = self.item_centroids[item2cluster]  # [B, e]
        pos_score_item = torch.mul(norm_item_embeddings, item2centroids).sum(dim=1)
        pos_score_item = torch.exp(pos_score_item / self.ssl_temp)
        ttl_score_item = torch.matmul(norm_item_embeddings, self.item_centroids.transpose(0, 1))
        ttl_score_item = torch.exp(ttl_score_item / self.ssl_temp).sum(dim=1)
        proto_nce_loss_item = -torch.log(pos_score_item / ttl_score_item).sum()

        # 两个损失函数的加权和，权重是一个很小的数
        proto_nce_loss = self.proto_reg * (proto_nce_loss_user + proto_nce_loss_item)
        return proto_nce_loss

    # 向量的结构对比学习损失计算
    # 现向量是卷积若干层后的向量，前向量是卷积时的初始向量
    def ssl_layer_loss(self, current_embedding, previous_embedding, user, item):
        current_user_embeddings, current_item_embeddings = torch.split(current_embedding, [self.n_users, self.n_items])
        previous_user_embeddings_all, previous_item_embeddings_all = torch.split(previous_embedding,
                                                                                 [self.n_users, self.n_items])

        current_user_embeddings = current_user_embeddings[user]
        previous_user_embeddings = previous_user_embeddings_all[user]
        norm_user_emb1 = F.normalize(current_user_embeddings)
        norm_user_emb2 = F.normalize(previous_user_embeddings)
        norm_all_user_emb = F.normalize(previous_user_embeddings_all)
        # pos是正样本对应的对比分子，ttl是全部（total）的对比分母
        pos_score_user = torch.mul(norm_user_emb1, norm_user_emb2).sum(dim=1)
        ttl_score_user = torch.matmul(norm_user_emb1, norm_all_user_emb.transpose(0, 1))
        pos_score_user = torch.exp(pos_score_user / self.ssl_temp)
        ttl_score_user = torch.exp(ttl_score_user / self.ssl_temp).sum(dim=1)

        ssl_loss_user = -torch.log(pos_score_user / ttl_score_user).sum()

        current_item_embeddings = current_item_embeddings[item]
        previous_item_embeddings = previous_item_embeddings_all[item]
        norm_item_emb1 = F.normalize(current_item_embeddings)
        norm_item_emb2 = F.normalize(previous_item_embeddings)
        norm_all_item_emb = F.normalize(previous_item_embeddings_all)
        pos_score_item = torch.mul(norm_item_emb1, norm_item_emb2).sum(dim=1)
        ttl_score_item = torch.matmul(norm_item_emb1, norm_all_item_emb.transpose(0, 1))
        pos_score_item = torch.exp(pos_score_item / self.ssl_temp)
        ttl_score_item = torch.exp(ttl_score_item / self.ssl_temp).sum(dim=1)

        ssl_loss_item = -torch.log(pos_score_item / ttl_score_item).sum()

        # 最终加权求和
        ssl_loss = self.ssl_reg * (ssl_loss_user + self.alpha * ssl_loss_item)
        return ssl_loss

    # 基于增强图的卷积对比损失
    def ssl_graph_loss(self, aug_user_all_emb1, aug_item_all_emb1,
                       aug_user_all_emb2, aug_item_all_emb2, user, item):

        norm_all_user_emb = F.normalize(aug_user_all_emb2)
        aug_user_all_emb1 = aug_user_all_emb1[user]
        aug_user_all_emb2 = aug_user_all_emb2[user]
        norm_user_emb1 = F.normalize(aug_user_all_emb1)
        norm_user_emb2 = F.normalize(aug_user_all_emb2)
        # pos是正样本对应的对比分子，ttl是全部（total）的对比分母
        pos_score_user = torch.mul(norm_user_emb1, norm_user_emb2).sum(dim=1)
        ttl_score_user = torch.matmul(norm_user_emb1, norm_all_user_emb.transpose(0, 1))
        pos_score_user = torch.exp(pos_score_user / self.ssl_temp)
        ttl_score_user = torch.exp(ttl_score_user / self.ssl_temp).sum(dim=1)

        ssl_loss_user = -torch.log(pos_score_user / ttl_score_user).sum()

        aug_item_all_emb1 = aug_item_all_emb1[item]
        previous_item_embeddings = aug_item_all_emb2[item]
        norm_item_emb1 = F.normalize(aug_item_all_emb1)
        norm_item_emb2 = F.normalize(previous_item_embeddings)
        norm_all_item_emb = F.normalize(aug_item_all_emb2)
        pos_score_item = torch.mul(norm_item_emb1, norm_item_emb2).sum(dim=1)
        ttl_score_item = torch.matmul(norm_item_emb1, norm_all_item_emb.transpose(0, 1))
        pos_score_item = torch.exp(pos_score_item / self.ssl_temp)
        ttl_score_item = torch.exp(ttl_score_item / self.ssl_temp).sum(dim=1)

        ssl_loss_item = -torch.log(pos_score_item / ttl_score_item).sum()

        # 最终加权求和
        ssl_loss = self.ssl_reg * (ssl_loss_user + self.alpha * ssl_loss_item)
        return ssl_loss

    # 计算增强图卷积前后的对比损失
    def aug_ssl_layer_loss(self, aug_user_all_emb, aug_item_all_emb, previous_embedding, user, item):
        previous_user_embeddings_all, previous_item_embeddings_all = torch.split(previous_embedding,
                                                                                 [self.n_users, self.n_items])

        aug_user_all_emb = aug_user_all_emb[user]
        previous_user_embeddings = previous_user_embeddings_all[user]
        norm_user_emb1 = F.normalize(aug_user_all_emb)
        norm_user_emb2 = F.normalize(previous_user_embeddings)
        norm_all_user_emb = F.normalize(previous_user_embeddings_all)
        # pos是正样本对应的对比分子，ttl是全部（total）的对比分母
        pos_score_user = torch.mul(norm_user_emb1, norm_user_emb2).sum(dim=1)
        ttl_score_user = torch.matmul(norm_user_emb1, norm_all_user_emb.transpose(0, 1))
        pos_score_user = torch.exp(pos_score_user / self.ssl_temp)
        ttl_score_user = torch.exp(ttl_score_user / self.ssl_temp).sum(dim=1)

        ssl_loss_user = -torch.log(pos_score_user / ttl_score_user).sum()

        aug_item_all_emb = aug_item_all_emb[item]
        previous_item_embeddings = previous_item_embeddings_all[item]
        norm_item_emb1 = F.normalize(aug_item_all_emb)
        norm_item_emb2 = F.normalize(previous_item_embeddings)
        norm_all_item_emb = F.normalize(previous_item_embeddings_all)
        pos_score_item = torch.mul(norm_item_emb1, norm_item_emb2).sum(dim=1)
        ttl_score_item = torch.matmul(norm_item_emb1, norm_all_item_emb.transpose(0, 1))
        pos_score_item = torch.exp(pos_score_item / self.ssl_temp)
        ttl_score_item = torch.exp(ttl_score_item / self.ssl_temp).sum(dim=1)

        ssl_loss_item = -torch.log(pos_score_item / ttl_score_item).sum()

        # 最终加权求和
        ssl_loss = self.ssl_reg * (ssl_loss_user + self.alpha * ssl_loss_item)
        return ssl_loss

    # 串起了训练的主流程
    def calculate_loss(self, interaction):
        # print("interaction", interaction)
        # clear the storage variable when training
        if self.restore_user_e is not None or self.restore_item_e is not None:
            self.restore_user_e, self.restore_item_e = None, None

        user = interaction[self.USER_ID]
        pos_item = interaction[self.ITEM_ID]
        # print("pos_item:", pos_item.shape, pos_item)
        neg_item = interaction[self.NEG_ITEM_ID]
        # print("neg_item:", neg_item.shape, neg_item)

        # 最终的用户项目向量，以及卷积过程中的全部向量组
        user_all_embeddings, item_all_embeddings, embeddings_list = self.forward(self.norm_adj_mat)

        # 取出用户项目的初始向量
        # center_embedding = embeddings_list[0]
        # 环境向量，取卷过几层后的向量
        # context_embedding = embeddings_list[self.hyper_layers * 2]

        # 将运算，得到两份增强图卷积向量
        aug_user_all_emb1, aug_item_all_emb1, _ = self.forward(self.aug_norm_mat1)
        aug_user_all_emb2, aug_item_all_emb2, _ = self.forward(self.aug_norm_mat2)
        # 对提取的当前批次的两份卷积向量计算对比损失
        aug_graph_loss = self.ssl_graph_loss(aug_user_all_emb1, aug_item_all_emb1,
                                             aug_user_all_emb2, aug_item_all_emb2, user, pos_item)

        # 计算初始向量和结构相邻向量的对比损失
        # ssl_loss = self.ssl_layer_loss(context_embedding, center_embedding, user, pos_item)
        # 计算初始向量和语义对比聚类向量的对比损失
        # proto_loss = self.ProtoNCE_loss(center_embedding, user, pos_item)
        # 计算初始向量和卷积后的语义对比聚类向量的对比损失
        # conv_loss = self.ConvNCE_loss(center_embedding, user, pos_item)

        # 计算两张增强图和卷积后的语义对比聚类向量的对比损失
        aug1_cv_loss = self.AugConvNCE_loss(aug_user_all_emb1, aug_item_all_emb1, user, pos_item, 1)
        aug2_cv_loss = self.AugConvNCE_loss(aug_user_all_emb2, aug_item_all_emb2, user, pos_item, 2)

        # 计算增强卷积向量前后对比损失
        # aug_con_loss1 = self.aug_ssl_layer_loss(aug_user_all_emb1, aug_item_all_emb1,
        #                                         center_embedding, user, pos_item)
        # aug_con_loss2 = self.aug_ssl_layer_loss(aug_user_all_emb2, aug_item_all_emb2,
        #                                         center_embedding, user, pos_item)

        u_embeddings = user_all_embeddings[user]
        pos_embeddings = item_all_embeddings[pos_item]
        neg_embeddings = item_all_embeddings[neg_item]

        # 在与训练后计算非对比的普通损失，BPR Loss
        pos_scores = torch.mul(u_embeddings, pos_embeddings).sum(dim=1)
        neg_scores = torch.mul(u_embeddings, neg_embeddings).sum(dim=1)

        mf_loss = self.mf_loss(pos_scores, neg_scores)

        # 最后计算L2正则化损失
        u_ego_embeddings = self.user_embedding(user)
        pos_ego_embeddings = self.item_embedding(pos_item)
        neg_ego_embeddings = self.item_embedding(neg_item)
        reg_loss = self.reg_loss(u_ego_embeddings, pos_ego_embeddings, neg_ego_embeddings)

        # 普通损失加起来，对比损失单独返回
        return mf_loss + self.reg_weight * reg_loss, aug_graph_loss, aug1_cv_loss + aug2_cv_loss

    def predict(self, interaction):
        # 基于矩阵乘法与求和实现分数预测
        user = interaction[self.USER_ID]
        item = interaction[self.ITEM_ID]

        user_all_embeddings, item_all_embeddings, _ = self.forward(self.norm_adj_mat)

        u_embeddings = user_all_embeddings[user]
        i_embeddings = item_all_embeddings[item]
        scores = torch.mul(u_embeddings, i_embeddings).sum(dim=1)
        return scores

    def full_sort_predict(self, interaction):
        # 获得指定用户的评分
        user = interaction[self.USER_ID]
        if self.restore_user_e is None or self.restore_item_e is None:
            self.restore_user_e, self.restore_item_e, embedding_list = self.forward(self.norm_adj_mat)
        # get user embedding from storage variable
        u_embeddings = self.restore_user_e[user]

        # dot with all item embedding to accelerate
        scores = torch.matmul(u_embeddings, self.restore_item_e.transpose(0, 1))

        return scores.view(-1)
